"""
数据库连接和会话管理
"""
from sqlalchemy import create_engine
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import declarative_base, sessionmaker
from sqlalchemy.pool import QueuePool, AsyncAdaptedQueuePool
from app.config import settings

# SQLAlchemy基础类
Base = declarative_base()

# 异步数据库引擎 - 使用AsyncAdaptedQueuePool
async_engine = create_async_engine(
    settings.database_url,
    poolclass=AsyncAdaptedQueuePool,
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,  # 连接前检查连接有效性
    pool_recycle=3600,   # 1小时后回收连接
    pool_reset_on_return='commit',  # 返回连接池时重置连接
    connect_args={
        "connect_timeout": 60,  # aiomysql支持的连接超时
        "charset": "utf8mb4",
        "autocommit": False,
    },
    echo=False,
    future=True
)

# 同步数据库引擎 - 使用QueuePool
sync_engine = create_engine(
    settings.sync_database_url,
    poolclass=QueuePool,
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,
    pool_recycle=3600,
    pool_reset_on_return='commit',  # 返回连接池时重置连接
    connect_args={
        "connect_timeout": 60,
        "read_timeout": 300,     # 读取超时（5分钟，适合大内容插入）
        "write_timeout": 300,     # 写入超时（5分钟）
        "charset": "utf8mb4",
        "autocommit": False,
    },
    echo=False,
    future=True
)

# 异步会话工厂
AsyncSessionLocal = sessionmaker(
    bind=async_engine,
    class_=AsyncSession,
    expire_on_commit=False
)

# 同步会话工厂
SessionLocal = sessionmaker(
    bind=sync_engine,
    autocommit=False,
    autoflush=False
)


async def get_async_db():
    """获取异步数据库会话"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


def get_db():
    """获取同步数据库会话"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


async def create_tables():
    """创建数据库表"""
    async with async_engine.begin() as conn:
        # 导入所有模型以确保表被创建
        from app.models.theme import Theme
        from app.models.article import Article
        from app.models.media import MediaAsset
        
        await conn.run_sync(Base.metadata.create_all)


async def close_db_connections():
    """关闭数据库连接"""
    try:
        # 先关闭异步引擎
        if async_engine:
            await async_engine.dispose()
    except Exception as e:
        print(f"关闭异步数据库连接时出错: {e}")
    
    try:
        # 再关闭同步引擎
        if sync_engine:
            sync_engine.dispose()
    except Exception as e:
        print(f"关闭同步数据库连接时出错: {e}")